import unittest
from Mixer import *

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')


class test(unittest.TestCase):

    @classmethod
    def setUpClass(cls) -> None:
        cls.sequence_length = (IMG_SHAPE[0] * IMG_SHAPE[1]) // (PATCH_SHAPE[0] * PATCH_SHAPE[1])
        cls.test_org_img = torch.randn(BTACH_SIZE, IMG_CHANNEL_NUM, *IMG_SHAPE)
        cls.test_embedding_img = torch.randn(BTACH_SIZE, cls.sequence_length, HIDDEN_CHANNEL_NUM)

    def test_mlp_block(self):
        mlp_block_ds = MLP_block(self.sequence_length, DS)
        mlp_block_dc = MLP_block(HIDDEN_CHANNEL_NUM, DC)
        self.assertEqual(
            self.test_embedding_img.shape,
            mlp_block_dc(self.test_embedding_img).shape
        )
        self.assertEqual(
            self.test_embedding_img.transpose(-2, -1).shape,
            mlp_block_ds(self.test_embedding_img.transpose(-2, -1)).shape
        )
        print(mlp_block_ds(self.test_embedding_img.transpose(-2, -1)).shape)
        print(mlp_block_dc(self.test_embedding_img).shape)

    def test_mixer_layer(self):
        mixer_layer = Mixer_layer((self.sequence_length, HIDDEN_CHANNEL_NUM), DS, DC)
        self.assertEqual(
            self.test_embedding_img.shape,
            mixer_layer(self.test_embedding_img).shape
        )
        print(mixer_layer(self.test_embedding_img).shape)

    def test_MLPmixer(self):
        mlp_mixer = MLP_Mixer(IMG_SHAPE, IMG_CHANNEL_NUM, PATCH_SHAPE, HIDDEN_CHANNEL_NUM, MIXER_LAYER_NUM, DS, DC,
                              OUTPUT_CLASS_NUM)
        self.assertEqual(
            [BTACH_SIZE, OUTPUT_CLASS_NUM],
            list(mlp_mixer(self.test_org_img).shape)
        )

    def test_print_net(self):
        from torchsummary import summary
        mlp_mixer = MLP_Mixer(IMG_SHAPE, IMG_CHANNEL_NUM, PATCH_SHAPE, HIDDEN_CHANNEL_NUM, MIXER_LAYER_NUM, DS, DC,
                              OUTPUT_CLASS_NUM)
        summary(mlp_mixer, input_size=(3, 224, 224), device='cpu')

        # stat(mlp_mixer,(3,224,224))

    def test_info(self):
        mlp_mixer = MLP_Mixer(IMG_SHAPE, IMG_CHANNEL_NUM, PATCH_SHAPE, HIDDEN_CHANNEL_NUM, MIXER_LAYER_NUM, DS, DC,
                              OUTPUT_CLASS_NUM)
        # print(mlp_mixer)
        for name, i in mlp_mixer.named_parameters():
            print(name, ':', i.data.shape)


if __name__ == '__main__':
    unittest.main()
